import json
import codecs
import random

def build_train_data_for_bge(path="zhuangzi.json", negtive_sample_rate=2):
    """
    这里不应该用句对句数据训练，因为会泄露，但是可以用段对段数据训练
    :param path:
    :param negtive_sample_rate:
    :return:
    """

    def merge_paragraph(paragraph: dict, for_classic: bool):
        if for_classic:
            return "".join([p["text"] for p in paragraph["comments"]])
        author_text = {}
        for comment in paragraph["comments"]:
            for author in comment:
                if author not in author_text:
                    author_text[author] = comment[author]
                else:
                    author_text[author] += "。"+comment[author]
        return author_text

    paragraphs = json.load(open(path))
    with open("zhuangzi_bge_train.jsonl", "w") as writer:
        for paragraph in paragraphs:
            s1 = merge_paragraph(paragraph, True)
            author_text = merge_paragraph(paragraph, False)
            for s2 in author_text.values():
                neg_paragraphs = random.sample(paragraphs, negtive_sample_rate)
                line_data = {"query": s1, "pos": [s2], "neg": [random.choice(list(merge_paragraph(p, False).values())) for p in neg_paragraphs]}
                writer.write(json.dumps(line_data, ensure_ascii=False) + "\n")
                line_data = {"query": s2, "pos": [s1], "neg": [merge_paragraph(p, True) for p in neg_paragraphs]}
                writer.write(json.dumps(line_data, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    build_train_data_for_bge()